from typing import Dict, List

from ..shard.placement_types import DTensorSpec,Placement,Replicate,PlacementStrategy,Shard,_Partial,ShardingStrategy
from .strategy_generator import StrategyGenerator

__all__ = ['PlaceholderGenerator']


class PlaceholderGenerator(StrategyGenerator):
    """
    PlaceholderGenerator is a generic class to generate strategies for placeholder node.
    删除了optional
    """

    def collate_strategies(self) -> List[ShardingStrategy]:
        dim_size = len(self.op_data['output'].data.shape)
        # sharding_specs_data = self.get_sharding_specs(sharding_specs_mapping)
        # sharding strategy bookkeeping
        strategy_list = []
        # 1、复制
        name = "() -> R"
        output_spec = DTensorSpec(mesh = self.device_mesh, placements=(Replicate(),),tensor_meta=self.op_data['output'].data)
        sharding_specs = PlacementStrategy(output_specs=output_spec)
        strategy_list.append(ShardingStrategy(name=name,sharding_specs=sharding_specs,compute_cost=0))
        if dim_size >= 1:
            name = "() -> S(0)"
            output_spec = DTensorSpec(mesh = self.device_mesh, placements=(Shard(0),),tensor_meta=self.op_data['output'].data)
            sharding_specs = PlacementStrategy(output_specs=output_spec)
            strategy_list.append(ShardingStrategy(name=name,sharding_specs=sharding_specs,compute_cost=0))
        if dim_size >= 2:
            name = "() -> S(1)"
            output_spec = DTensorSpec(mesh = self.device_mesh, placements=(Shard(1),),tensor_meta=self.op_data['output'].data)
            sharding_specs = PlacementStrategy(output_specs=output_spec)
            strategy_list.append(ShardingStrategy(name=name,sharding_specs=sharding_specs,compute_cost=0))
        # elif dim_size == 3:
        #     string = "mkn->mkn"
        return strategy_list
